import math
import os
import cv2

LABEL_DATA_PATH = "/disk527/DataDisk/a804_cbf/datasets/lunar_crater_database_robbins_2018_bundle/data/lunar_crater_database_robbins_2018.csv"

SUBFIG_DIR = "/disk527/sdb1/a804_cbf/datasets/lunar_crater"

LATITUDE_RANGE = 0.5084  # degree
LONGITUDE_RANGE = 9.1035  # degree
X_PIXEL_RANGE = 12060
Y_PIXEL_RANGE = 11660
MOON_RADIUS = 1737.4  # km
BATCH_SIZE = 1200


def grey_image_div(output_dir, batch_size=1200):
    if not os.path.exists(f"{output_dir}/images"):
        os.makedirs(f"{output_dir}/images")
    IMAGE_PATH = "/disk527/DataDisk/a804_cbf/datasets/lunar_crater/Lunar_LROnac_Haworth_sfs-dem_1m_v3.tif"
    image = cv2.imread(IMAGE_PATH, cv2.IMREAD_UNCHANGED)
    num_x = image.shape[1] // batch_size
    num_y = image.shape[0] // batch_size
    for i in range(num_x):
        for j in range(num_y):
            sub_image = image[
                j * batch_size : (j + 1) * batch_size,
                i * batch_size : (i + 1) * batch_size,
            ]
            cv2.imwrite(f"{output_dir}/images/{i}_{j}.tif", sub_image)


def mercator_projection_R(r, lat):
    R_x = math.degrees(r / MOON_RADIUS) * X_PIXEL_RANGE / LATITUDE_RANGE
    R_y = (
        math.degrees(r / MOON_RADIUS)
        * Y_PIXEL_RANGE
        / LONGITUDE_RANGE
        / math.cos(math.radians(lat))
    )
    return R_x, R_y


# Here is the definition of image coordinate system
# +-------> y
# | (0, 0)
# |
# |
# |
# v
# x
def mercator_projection(lat, lon):
    """
    Warning ! The unit of latitude and longitude is degree.
    """
    # 这里要加上一半的图像宽度，这是因为标注数据集中的角度是从-60度到+60度，起始位置不为零，因此需要修正。
    x = (-86.59046214259241 - lat) * X_PIXEL_RANGE / LATITUDE_RANGE
    # 这里不加一半的图像高度，这是因为标注数据集中的角度是从0到360度的，但是，图像的经度是从-180到+180度的，因此需要修正。
    y = (-16.624216585089282 - lon) * Y_PIXEL_RANGE / LONGITUDE_RANGE
    return x, y


def labels(line, header):
    lat, lon = float(line[1]), float(line[2])
    lon = lon if lon < 180 else lon - 360
    if (
        lat < -87.09890342731742
        or lat > -86.59046214259241
        or lon < -25.72774505919917
        or lon > -16.624216585089282
    ):
        raise ValueError("Out of range.")
    # for circle projection
    temp_1, temp_2 = mercator_projection(lat, lon)
    R_x, R_y = mercator_projection_R(float(line[5]), lat)
    top_left = (temp_1 - R_x, temp_2 - R_y)
    bottom_right = (temp_1 + R_x, temp_2 + R_y)
    if (
        top_left[0] < 0
        or top_left[1] < 0
        or bottom_right[0] > X_PIXEL_RANGE
        or bottom_right[1] > Y_PIXEL_RANGE
    ):
        raise ValueError("Out of range.")
    x_folder = int(temp_1) // BATCH_SIZE
    y_folder = int(temp_2) // BATCH_SIZE
    x = temp_1 - BATCH_SIZE * x_folder
    y = temp_2 - BATCH_SIZE * y_folder
    cir_labels = {"Y_FOLDER": y_folder, "X_FOLDER": x_folder, "Y_SUB": y, "X_SUB": x}

    for field, val in zip(header, line):
        if "CIRC" in field:
            cir_labels = dict(**{field: val}, **cir_labels)
    return cir_labels


if __name__ == "__main__":
    # 从标注文件中获取某一行
    with open(LABEL_DATA_PATH, "r") as f:
        # 有部分行只有圆的数据，没有椭圆的数据
        # 流式读取，第一行是标题，并使用标题的字段名作为数据变量名
        header = f.readline()
        header = header.strip().split(",")
        while True:
            # 读取第一行数据
            line_ = f.readline()
            if not line_:
                break
            line = line_.strip().split(",")
            try:
                cir_labels = labels(line, header)
            except:
                continue
            cir_x_folder = cir_labels["X_FOLDER"]
            cir_y_folder = cir_labels["Y_FOLDER"]
            if (
                cir_x_folder > 7
                or cir_x_folder < 0
                or cir_y_folder > 7
                or cir_y_folder < 0
            ):
                continue
            if not os.path.exists(f"{SUBFIG_DIR}/labels/circle/{cir_x_folder}"):
                os.makedirs(f"{SUBFIG_DIR}/labels/circle/{cir_x_folder}")
            with open(
                f"{SUBFIG_DIR}/labels/circle/{cir_x_folder}/lbl_{cir_x_folder}_{cir_y_folder}_{BATCH_SIZE}.txt",
                "a",
            ) as writer:
                temp = list(cir_labels.values())
                temp.reverse()
                writer.write(",".join(map(lambda x: str(x), temp)) + "\n")
                writer.flush()
    grey_image_div("/disk527/sdb1/a804_cbf/datasets/harworth", batch_size=256)
